import rosbag
import numpy as np
import sensor_msgs.point_cloud2 as pc2
import open3d as o3d
import cv2
import tqdm
from cv_bridge import CvBridge
from tictoc import TicToc
import csv
import os
import json
import time
import argparse
import multiprocessing
from cyw_devkit.core.visualizer.zwh_viewer import tranform_pt_pad
from my_io import write_points, copy_file, files_filter, read_points

now = time.time()

# 防止多进程时间戳重复，用进程id拉开一小时 3600*1000ms间隔
def get_ms(msg, process_id=0):
    if msg.header.stamp.secs == 0 and msg.header.stamp.nsecs == 0:
        return int(msg.time_us / 1000)
    return (msg.header.stamp.secs + int(now)) * 1000 + int(
        msg.header.stamp.nsecs / (10 ** 6) + process_id * 3600 * 1000)

def main(args):
    # 统计耗时
    cost = TicToc("ros解析")
    args.rosbag_dir = os.path.join(args.work_space, 'bags')
    assert os.path.exists(args.rosbag_dir)
    files = files_filter(args.rosbag_dir, pos_fix="bag", split='.')
    process_size = len(files)
    manager = multiprocessing.Manager()
    if process_size > 1:
        pool = multiprocessing.Pool(process_size)
        counter_list = manager.list()
        for idx in range(process_size):
            pool.apply_async(main_worker, args=(files[idx], args.work_space))
        pool.close()
        pool.join()
    else:
        main_worker(files[0], args.work_space)

    print("---------------------------------------------------------")
    print("处理完成: {}".format(files))
    cost.toc()
    print("---------------------------------------------------------")


def main_worker(bag, work_space):
    rosbag_dir = os.path.join(work_space, 'bags')
    data_path = os.path.join(work_space, 'datasets')

    rosbag_path = os.path.join(rosbag_dir, bag)
    dir_name = os.path.join(data_path, '__'+bag.split('.')[0])
    print("开始解析<{}>......".format(rosbag_path))

    bag = rosbag.Bag(rosbag_path)
    #
    lidar_topics = ["/livox/lidar"]
    lidar_dim=4
    print("保存在<{}>".format(dir_name))

    # 新建目录
    new_dir = []
    ffarther_dir = []
    farther_dir = []
    child_dir = []
    ffarther_dir.append(dir_name)
    new_dir = new_dir + ffarther_dir
    farther_dir.append(os.path.join(dir_name, "lidar_bin"))
    farther_dir.append(os.path.join(dir_name, "calib"))
    farther_dir.append(os.path.join(dir_name, "localization"))
    new_dir = new_dir + farther_dir
    new_dir = new_dir + child_dir
    for dir in new_dir:
        if not os.path.exists(dir):
            print("新建目录<{}>".format(dir))
            os.makedirs(dir, exist_ok=True)

    if os.path.exists(os.path.join(rosbag_dir, 'calib.json')):
        copy_file(os.path.join(rosbag_dir, 'calib.json'), os.path.join(dir_name, 'calib'))
        print("拷贝 calib.json , from {} to {}".format(os.path.join(rosbag_dir, 'calib'),
                                                       os.path.join(dir_name, 'calib')))

    calib_json = os.path.join(dir_name, 'calib', 'calib.json')
    static_transfoms = dict()
    with open(calib_json, 'r') as f:
        json_data = json.load(f)
        for key, val in json_data.items():
            if isinstance(val, dict) and 'transforms' in val.keys():
                for skey, sval in val['transforms'].items():
                    static_transfoms[key] = {skey: np.array(sval)}

    # 创建或打开文件
    csvfile = open(os.path.join(dir_name, "localization", 'localization.csv'), mode='w', newline='')
    # 标题列表
    fieldnames = ['stamp(ms)', 'latitude', 'longitude', 'altitude', 'mercator_x', 'mercator_y', 'mercator_z',
                  'vel_x', 'vel_y', 'vel_z', 'accel_x', 'accel_y', 'accel_z',
                  'roll', 'pitch', 'yaw', 'angular_x', 'angular_y', 'angular_z']
    # 创建 DictWriter 对象
    write = csv.DictWriter(csvfile, fieldnames=fieldnames)
    # 写入表头
    write.writeheader()

    for topic, msg, t in bag:
        if topic in lidar_topics:
            lidar = pc2.read_points(msg)
            points = np.array(list(lidar))
            points = points[(points[:,0] > 0.1) & (points[:,0] < 50.0)]
            points = tranform_pt_pad(points,tfs=[static_transfoms['lidar3']['base_link']])
            # # 创建一个PointCloud对象
            # pcd = o3d.geometry.PointCloud()
            # # 将随机数转换成PointCloud点数据
            # pcd.points = o3d.utility.Vector3dVector(points[:, :3])
            # # 将PointCloud点数据保存成pcd文件，格式为assii文本格式 比如用于lc标定任务
            # o3d.io.write_point_cloud(os.path.join(dir_name,'lidar', topic2path(topic), str(get_ms(msg)) + ".pcd"), pcd,
            #                          write_ascii=True)
            assert points.shape[1] >= 4
            write_points(points[:, :lidar_dim],
                         os.path.join(dir_name, 'lidar_bin', str(get_ms(msg)) + ".bin"))

            write.writerow({'stamp(ms)': str(get_ms(msg)) + '\t', 'latitude': 0.0, 'longitude': 0.0,
                            'altitude': 0.0, 'mercator_x': 0.0, 'mercator_y': 0.0,
                            'mercator_z': 0.0, 'vel_x': 0.0, 'vel_y': 0.0, 'vel_z': 0.0,
                            'accel_x': 0.0,
                            'accel_y': 0.0, 'accel_z': 0.0, 'roll': 0.0, 'pitch': 0.0,
                            'yaw': 0.0,
                            'angular_x': 0.0, 'angular_y': 0.0,
                            'angular_z': 0.0})

    for dir in child_dir:
        print("文件夹:<{}> 包含文件<{}>个".format(dir.split('/')[-1], len(os.listdir(dir))))

    bag.close()

    # 关键帧
    output_paths_samples = os.path.join(dir_name, "samples")
    if not os.path.exists(output_paths_samples):
        os.mkdir(output_paths_samples)

    lidar_len = len(os.listdir(os.path.join(dir_name, 'lidar_bin')))
    keep_data_idx = np.arange(start=0,stop=lidar_len)[::3]
    key_info = {
        'key_index': keep_data_idx.tolist()
    }

    # 存关键帧信息, 写入JSON文件
    with open(os.path.join(output_paths_samples, 'samples.json'), 'w', encoding='utf-8') as file:
        json.dump(key_info, file, ensure_ascii=False, indent=4)

    print('全部帧：{}, 关键帧数据:{}'.format(lidar_len, len(keep_data_idx)))

    # 新建readme说明文件
    # 创建或打开文件
    csvfile = open(os.path.join(dir_name, 'README.csv'), mode='w', newline='')
    # 标题列表
    fieldnames = ['Item', 'Details']
    # 创建 DictWriter 对象
    write = csv.DictWriter(csvfile, fieldnames=fieldnames)
    # 写入表头
    write.writeheader()
    write.writerow({'Item': 'Bag_Name', 'Details': dir_name.split('/')[-1]})
    write.writerow({'Item': 'With_Intensity', 'Details': True})
    write.writerow({'Item': 'Description', 'Details': '<you can note something here>'})
    write.writerow({'Item': 'SyncParas',
                    'Details': 'max_continuous:{} key_dis:{}'.format(3, -1)})
    write.writerow({'Item': 'Frames_Size', 'Details': lidar_len})
    write.writerow({'Item': 'Samples_Size', 'Details': len(keep_data_idx)})
    write.writerow({'Item': 'Topics_Size', 'Details': 1})
    write.writerow({'Item': 'Topics', 'Details': "livox/lidar"})
    write.writerow({'Item': 'Split(train:val:test)', 'Details': "Waitting For Split"})
    write.writerow({'Item': 'Mission', 'Details': ['det', 'lidar']})
    write.writerow({'Item': 'Version', 'Details': "livox"})

    # 待标注
    tmp_dir = "anno_" + os.path.basename(dir_name)
    annotate_dir = os.path.join(work_space, 'annotate', tmp_dir)

    tmp_anotate = os.path.join(annotate_dir, 'lidar')
    label_dir = os.path.join(annotate_dir, 'label')
    os.makedirs(tmp_anotate, exist_ok=True)
    os.makedirs(label_dir, exist_ok=True)

    lidars = sorted(os.listdir(os.path.join(dir_name, 'lidar_bin')))
    for idx in keep_data_idx:
        pc = read_points(os.path.join(dir_name,'lidar_bin',lidars[idx]), dim=lidar_dim)
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(pc[:, :3])
        pcd_file = os.path.join(tmp_anotate, lidars[idx].replace('.bin', '.pcd'))
        o3d.io.write_point_cloud(pcd_file, pcd,
                                 write_ascii=True)  # save xyz.pcd
    pass


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--work-space',
                        default='/home/adt/bags/livox',
                        help='按照格式整理好')
    args = parser.parse_args()

    main(args)
